from abc import abstractmethod
from collections import deque
from xml.etree import ElementTree as ET
from typing import Iterable, Optional, Union, Any
from torch._tensor import Tensor
from tqdm import tqdm
from feasytools import ArgChecker
import bisect
from utils.draw import plot_results
from utils.data_loader import *
from utils.tools import metric
from models import Pyraformer, WorkerArguments, get_model_save_dir, create_model
from preproc import get_trips
import copy


class _TPredFrontEnd:
    def __init__(self, model_pars: 'Union[WorkerArguments, dict[str, Any]]'):
        if isinstance(model_pars, dict):
            model_pars = WorkerArguments(model_pars)
        mdir = get_model_save_dir(model_pars)
        self.save_root = mdir + "results_rolling"
        self.__pars = model_pars
        self.__buf_tl: 'list[int]' = []
        self.__buf_dt: 'list[torch.Tensor]' = []
        self._tl: 'deque[int]' = deque()
        self._dt: 'deque[torch.Tensor]' = deque()
        self.__trips_tl: 'list[int]' = []
        self.__trips_dt: 'list[torch.Tensor]' = []

    @property
    def ilen(self) -> int:
        """输入的时间跨度，单位分钟"""
        return self.__pars.ilen

    @property
    def olen(self) -> int:
        """输出的时间跨度，单位分钟"""
        return self.__pars.olen

    @property
    def isteps(self) -> int:
        """输入的序列长度"""
        return self.__pars.input_size

    @property
    def osteps(self) -> int:
        """输出的序列长度"""
        return self.__pars.predict_step

    @property
    def sr(self) -> int:
        """序列中数据点的时间间隔，单位分钟"""
        return self.__pars.sr

    @property
    def dev(self) -> torch.device:
        return self.__pars.device

    def __add(self, time: int, data: torch.Tensor):
        self._tl.append(time)
        self._dt.append(data)
        while len(self._tl) > self.isteps:
            self._tl.popleft()
            self._dt.popleft()

    def __addrange(self, time: Iterable[int], data: Iterable[torch.Tensor]):
        self._tl.extend(time)
        self._dt.extend(data)
        while len(self._tl) > self.isteps:
            self._tl.popleft()
            self._dt.popleft()

    def set_trips(self, data: ET.Element) -> None:
        """设置行程信息"""
        self.__trips_head, trips = get_trips(data)
        self.__trips_tl = list(trips[:, 0])
        self.__trips_dt = [torch.tensor(x) for x in trips[:, 1:]]

    def add_data(self, time: int, dataraw: Iterable[float]) -> bool:
        """
        添加数据，并移除陈旧数据
            time: 以秒为单位的时间戳
            dataraw: 不包含行程的数据(对于无V2G的情况, 应该是价格与负荷的拼接; 对于有V2G的情况, 应该是购电价格、售电价格和负荷的拼接)
        返回:
            已有数据是否足以进行一次预测
        """
        if len(self.__buf_tl) > 0: assert self.__buf_tl[-1] < time, "时间戳必须递增"
        assert len(self.__trips_tl) > 0 and len(self.__trips_dt) > 0, "行程信息未设置"
        data = torch.tensor(list(dataraw))
        if len(self.__buf_tl) == 0:
            self.__buf_tl.append(time)
            self.__buf_dt.append(data)
            return False
        if time - self.__buf_tl[0] >= self.sr * 60:
            tm = self.__buf_tl[0] // 60
            pos = bisect.bisect_left(self.__trips_tl, tm)
            assert self.__trips_tl[pos] == tm, f"未查找到{tm}时间的行程信息"
            trips_data = self.__trips_dt[pos]
            other_data = torch.stack(self.__buf_dt).mean(0)
            self.__add(tm, torch.cat((trips_data, other_data), dim=0))
            self.__buf_tl = []
            self.__buf_dt = []
        self.__buf_tl.append(time)
        self.__buf_dt.append(data)
        return len(self._tl) == self.isteps

    def load_data(self, dset: TrafficData, clear: bool = False):
        """
        从单周交通数据集中导入足够预测一次的数据
            dset: 交通数据集
            clear: 是否清空之前的数据(包括行程信息、负荷信息、V2G信息和价格信息等)
        """
        if clear:
            self.__buf_dt.clear()
            self.__buf_tl.clear()
            self._tl.clear()
            self._dt.clear()
        self.__addrange(
            dset.org_stamp[: self.isteps], torch.tensor(dset.data_x[: self.isteps])
        )

    @abstractmethod
    def predict(self) -> 'tuple[list[int], torch.Tensor]':
        raise NotImplementedError
    
    def iterative_predict(self, dset: TrafficData, pbar:Optional[tqdm] = None) -> 'tuple[list[int], np.ndarray]':
        """迭代预测数据，返回时间戳(以分钟计)和预测数据"""
        retT: 'Optional[list[int]]' = None
        retD: 'Optional[torch.Tensor]' = None
        while True:
            time0, data0 = self.predict()
            osteps, out_fea = data0.shape
            if retT is None or retD is None:
                retT = time0
                retD = data0
            else:
                retT.extend(time0)
                retD = torch.cat((retD, data0))
            newdata = dset.get_data_x(time0[0], time0[-1] + 1)
            if newdata is None:
                break
            newdata[:, -out_fea:] = data0.cpu().detach().numpy()
            self.__addrange(time0, torch.tensor(newdata))
            if pbar: pbar.update(1)
        return retT, retD.cpu().detach().numpy()

    def iterative_predict_and_compare(self, dset_all: Dataset_Traffic):
        """迭代预测单周的数据，并绘图比较"""
        all_preds = []
        all_trues = []
        with tqdm(total=dset_all._slen*(len(dset_all.data[0].data_x)-self.isteps)//self.osteps,
                desc="Predicting", leave=True, mininterval=1) as pbar:
            for dset in dset_all.data:
                self.load_data(dset)
                _, preds = self.iterative_predict(dset,pbar)
                steps, out_fea = preds.shape
                if self.__pars.non_future:
                    trues = dset.data_y[:steps, -out_fea:]
                else:
                    trues = dset.data_y[self.isteps : self.isteps + steps, -out_fea:]
                real_steps, _ = trues.shape
                all_preds.append(preds[:real_steps])
                all_trues.append(trues)
        os.makedirs(self.save_root, exist_ok=True)
        all_preds = np.stack(all_preds)
        all_preds = dset_all.scaler.inverse_transform(all_preds)
        all_trues = np.stack(all_trues)
        print("Metrics:", metric(all_preds, all_trues))
        plot_results(all_preds, all_trues, dset.cfg.datatype, self.save_root)

class _TPredBackEnd:
    @staticmethod
    def __tstamp(df_stamp: Iterable):
        return torch.tensor([(
            (x % 60 - 30) / 30.0,
            ((x // 60) % 24 - 12) / 12.0,
            (x // 1440 - 3) / 3.0,
        ) for x in df_stamp])
    
    def __init__(self, model_pars: 'Union[WorkerArguments, dict[str, Any]]', scaler: Optional[Scaler]):
        if isinstance(model_pars, dict):
            model_pars = WorkerArguments(model_pars)
        mdir = get_model_save_dir(model_pars)
        self.__model = create_model(model_pars)
        checkpoint = torch.load(mdir + "best_iter0.pth")["state_dict"]
        self.__model.load_state_dict(checkpoint)
        self.__pars = model_pars
        self.sr = self.__pars.sr
        self.isteps = self.__pars.input_size
        self.osteps = self.__pars.predict_step
        self.dev = self.__pars.device
        self.scaler = scaler
        
    @property
    def model(self) -> torch.nn.Module:
        return self.__model

    def predict(self, tl:'deque[int]', dt:'deque[torch.Tensor]') -> 'tuple[list[int], torch.Tensor]':
        """预测数据，返回时间戳(以分钟计)和预测数据"""
        assert len(tl) == self.isteps
        x: torch.Tensor = torch.stack(tuple(dt)).unsqueeze(0)
        if self.scaler:
            x = self.scaler.transform(x)
        x_mark: torch.Tensor = self.__tstamp(tl).unsqueeze(0)
        y: torch.Tensor = torch.zeros(self.osteps, self.__pars.dec_in).unsqueeze(0)
        if self.__pars.non_future:
            ty = [t * self.sr + tl[0] for t in range(self.osteps)]
        else:
            ty = [t * self.sr + tl[-1] for t in range(1, self.osteps + 1)]
        y_mark: torch.Tensor = self.__tstamp(ty).unsqueeze(0)
        x, x_mark, y, y_mark = map(
            lambda x: x.float().to(self.dev), (x, x_mark, y, y_mark)
        )
        if isinstance(self.model, Pyraformer) and self.__pars.decoder == "FC":
            predict_token = torch.zeros(x.size(0), 1, x.size(-1), device=x.device)
            x = torch.cat([x, predict_token], dim=1)
            x_mark = torch.cat([x_mark, y_mark[:, 0:1, :]], dim=1)

        pred: torch.Tensor = self.model(x, x_mark, y, y_mark, False)
        if self.__pars.scaler == "0max":
            pred = pred.clamp(-1, 1)
        if self.scaler:
            pred = self.scaler.inverse_transform(pred)
        return ty, pred.squeeze(0)

class TrafficPredictor(_TPredFrontEnd):
    def __init__(self, model_pars: 'Union[WorkerArguments, dict[str, Any]]', scaler: Optional[Scaler]):
        super().__init__(copy.deepcopy(model_pars))
        self.__be = _TPredBackEnd(model_pars, scaler)
    
    def predict(self) -> 'tuple[list[int], torch.Tensor]':
        return self.__be.predict(self._tl, self._dt)

if __name__ == "__main__":
    wa = WorkerArguments(ArgChecker())
    wa.inverse = True
    print(wa)
    ds = Dataset_Traffic(wa.path, wa, "val")
    predictor = TrafficPredictor(wa, None)
    predictor.iterative_predict_and_compare(ds)
